package cn.aulang.study.ai.client.config;

import io.micrometer.observation.ObservationRegistry;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
import org.springframework.ai.model.ollama.autoconfigure.OllamaChatProperties;
import org.springframework.ai.model.ollama.autoconfigure.OllamaInitializationProperties;
import org.springframework.ai.model.tool.DefaultToolExecutionEligibilityPredicate;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionEligibilityPredicate;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.ollama.management.PullModelStrategy;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.Objects;

/**
 * Tools服务
 *
 * @author Aulang
 */
@Configuration
public class ChatClientConfig {

    @Bean
    public OllamaQwenChatModel ollamaAlibabaChatModel(OllamaApi ollamaApi,
                                                      OllamaChatProperties chatProperties,
                                                      OllamaInitializationProperties initProperties,
                                                      ToolCallingManager toolCallingManager,
                                                      ObjectProvider<ObservationRegistry> observationRegistry,
                                                      ObjectProvider<ChatModelObservationConvention> observationConvention,
                                                      ObjectProvider<ToolExecutionEligibilityPredicate> ollamaToolExecutionEligibilityPredicate) {
        PullModelStrategy chatModelPullStrategy = initProperties.getChat().isInclude() ? initProperties.getPullModelStrategy() : PullModelStrategy.NEVER;
        OllamaQwenChatModel ollamaQwenChatModel = OllamaQwenChatModel
                .builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(chatProperties.getOptions())
                .toolCallingManager(toolCallingManager)
                .toolExecutionEligibilityPredicate(ollamaToolExecutionEligibilityPredicate.getIfUnique(DefaultToolExecutionEligibilityPredicate::new))
                .observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
                .modelManagementOptions(new ModelManagementOptions(chatModelPullStrategy, initProperties.getChat().getAdditionalModels(), initProperties.getTimeout(), initProperties.getMaxRetries()))
                .build();
        Objects.requireNonNull(ollamaQwenChatModel);
        observationConvention.ifAvailable(ollamaQwenChatModel::setObservationConvention);
        return ollamaQwenChatModel;
    }

    @Bean
    public ChatClient chatClient(@Autowired OllamaQwenChatModel ollamaQwenChatModel,
                                 @Autowired ToolCallbackProvider toolCallbackProvider) {
        return ChatClient.builder(ollamaQwenChatModel).defaultToolCallbacks(toolCallbackProvider).build();
    }
}
